from omegaconf import DictConfig
from lightning import Trainer
from lightning.pytorch.loggers import WandbLogger

from ecgcmr.imaging.img_dataset.ContrastiveImagingLightning import ContrastiveImageDataModule, ContrastiveImageDataModule_PyTorch
from ecgcmr.imaging.img_models.ImageSimCLR_ResNet3D import ImageEncoderSimCLR_ResNet3D
from ecgcmr.imaging.img_models.ImageSimCLR_ViT import ImageEncoderSimCLR_ViT
from ecgcmr.imaging.img_models.ImageBYOL_ViT import BYOLImageEncoder


def train_imaging_contrastive(
        cfg: DictConfig,
        wandb_logger: WandbLogger,
        save_dir: str,
        devices: int = 1
    ):

    if cfg.models.backbone == 'resnet':
        model = ImageEncoderSimCLR_ResNet3D(cfg=cfg, save_dir=save_dir)
    elif cfg.models.backbone == 'vit':
        model = ImageEncoderSimCLR_ViT(cfg=cfg, save_dir=save_dir)
    elif cfg.models.backbone == 'byol_vit':
        model = BYOLImageEncoder(cfg=cfg, save_dir=save_dir)

    wandb_logger.watch(model, log_graph=False)
    
    datamodule = ContrastiveImageDataModule(cfg=cfg)

    strategy = "ddp" if devices > 1 else "auto"

    trainer = Trainer(
        accelerator="gpu",
        devices=devices,
        strategy=strategy,
        precision="bf16-mixed",
        logger=wandb_logger,
        max_epochs=cfg.max_epochs,
        log_every_n_steps=cfg.log_every_n_steps,
        check_val_every_n_epoch=cfg.check_val_every_n_epoch,
        default_root_dir=save_dir,
        num_sanity_val_steps=0,
        profiler="simple",
        sync_batchnorm=True
    )
    
    trainer.fit(model=model, datamodule=datamodule)